# classfication.py
#
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

EPOCH = 120

n_data = torch.ones(2000, 2)    # shape = (2000,2)
x0 = torch.normal(2*n_data, 1)      # class0 x data (tensor), shape=(100, 2)
y0 = torch.zeros(2000)               # class0 y data (tensor), shape=(100, 1)
x1 = torch.normal(-2*n_data, 1)     # class1 x data (tensor), shape=(100, 2)
y1 = torch.ones(2000)                # class1 y data (tensor), shape=(100, 1)
# shape (200, 2) FloatTensor = 32-bit floating
x = torch.cat((x0, x1), 0).type(torch.FloatTensor)
# shape (200,) LongTensor = 64-bit integer
y = torch.cat((y0, y1), ).type(torch.LongTensor)


class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden)
        self.predict = torch.nn.Linear(n_hidden, n_output)

    def forward(self, x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        return x


net = Net(2, 20, 2)
print(net)
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
loss_func = torch.nn.CrossEntropyLoss()

loop_break = False

plt.ion()


for epoch in range(EPOCH):
    out = net(x)
    loss = loss_func(out, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 5 == 0:
        # plot and show learning process
        plt.cla()
        prediction = torch.max(out, 1)[1]
        pred_y = prediction.data.numpy()
        target_y = y.data.numpy()
        plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[
                    :, 1], c=pred_y, s=10, lw=0, cmap='RdYlGn')
        accuracy = float((pred_y == target_y).astype(
            int).sum()) / float(target_y.size)

        # loop_break = accuracy > 0.996

        plt.text(1.5, -4, 'Epoch=%i' %
                 epoch, fontdict={'size': 10, 'color': 'red'})
        plt.text(1.5, -5, 'Accuracy=%.5f' %
                 accuracy, fontdict={'size': 10, 'color': 'red'})

        plt.savefig(f'e:/temp/{int(epoch/5)}.png')
        plt.pause(0.5)

    if loop_break:
        break
# for epoch in range(EPOCH):
plt.ioff()
plt.show()
